#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <mpi.h>
#include "md5_parallel.h"

const uint32_t T[64] = {
	0xd76aa478,	0xe8c7b756,	0x242070db,	0xc1bdceee,	0xf57c0faf,	0x4787c62a,	0xa8304613,	0xfd469501,
	0x698098d8,	0x8b44f7af,	0xffff5bb1,	0x895cd7be,	0x6b901122,	0xfd987193,	0xa679438e,	0x49b40821,
	0xf61e2562,	0xc040b340,	0x265e5a51,	0xe9b6c7aa,	0xd62f105d,	0x2441453,	0xd8a1e681,	0xe7d3fbc8,
	0x21e1cde6,	0xc33707d6,	0xf4d50d87,	0x455a14ed,	0xa9e3e905,	0xfcefa3f8,	0x676f02d9,	0x8d2a4c8a,
	0xfffa3942,	0x8771f681,	0x6d9d6122,	0xfde5380c,	0xa4beea44,	0x4bdecfa9,	0xf6bb4b60,	0xbebfbc70,
	0x289b7ec6,	0xeaa127fa,	0xd4ef3085,	0x4881d05,	0xd9d4d039,	0xe6db99e5,	0x1fa27cf8,	0xc4ac5665,
	0xf4292244,	0x432aff97,	0xab9423a7,	0xfc93a039,	0x655b59c3,	0x8f0ccc92,	0xffeff47d,	0x85845dd1,
	0x6fa87e4f,	0xfe2ce6e0,	0xa3014314,	0x4e0811a1,	0xf7537e82,	0xbd3af235,	0x2ad7d2bb,	0xeb86d391
};

inline uint32_t swap_bytes(uint32_t in)
{
	return (in << 24) | ((in & 0xFF00) << 8) | ((in & 0xFF0000) >> 8) | (in >> 24);
}

uint32_t* pad_bits(uint8_t *message, uint64_t *length)
{
	uint8_t *new_message;
	uint32_t *out;
	uint64_t orig_length;
	int bytes_to_pad;
	int i;
	
	orig_length = *length;
	bytes_to_pad = (448 - ((int)orig_length << 3) % 512) >> 3;
	
	new_message = realloc(message, orig_length + bytes_to_pad + 8);
	out = calloc((orig_length+bytes_to_pad+8)/4,sizeof(uint32_t));
	
	new_message[orig_length] = 0x80;
		
	for(i = 0; i < orig_length+1; i+=4)
	{
		out[i/4] = new_message[i] | (new_message[i+1] << 8) | (new_message[i+2] << 16) | (new_message[i+3] << 24);
	}
	
	//out[*length/4] = 0x80;
	//memset(out+*length/4+1,0x00,bytes_to_pad - 4);
	
	*length += bytes_to_pad + 8;
	*length /= 4;
	
	out[*length - 2] = (orig_length*8) & 0xFFFFFFFF;
	out[*length - 1] = ((orig_length*8) & 0xFFFFFFFF00000000) >> 32;
		
	return out;
}

inline uint32_t rotate_left(uint32_t in, int n)
{
	return (in << n) | (in >> (32 - n));
}

inline uint32_t F(uint32_t X, uint32_t Y, uint32_t Z)
{
	return (X & Y) | ((~X) & Z);
}

inline uint32_t G(uint32_t X, uint32_t Y, uint32_t Z)
{
	return (X & Z) | (Y & (~Z));
}

inline uint32_t H(uint32_t X, uint32_t Y, uint32_t Z)
{
	return X ^ Y ^ Z;
}

inline uint32_t I(uint32_t X, uint32_t Y, uint32_t Z)
{
	return Y ^ (X | (~Z));
}

inline uint32_t round_1(uint32_t a, uint32_t b, uint32_t c, uint32_t d, uint32_t X_k, uint32_t s, uint32_t T_i)
{
	MPI_Request requests[6];
	MPI_Status statuses[6];
	uint32_t out;
	out = b + rotate_left((a + F(b,c,d) + X_k + T_i), s);
	
	//printf("out = %x, b = %x, c = %x, d = %x\n",out,b,c,d);
	MPI_Isend(&out,1,MPI_UNSIGNED,1,3,MPI_COMM_WORLD,requests);
	MPI_Isend(&b,1,MPI_UNSIGNED,1,2,MPI_COMM_WORLD,requests+1);
	MPI_Isend(&c,1,MPI_UNSIGNED,1,1,MPI_COMM_WORLD,requests+2);
	MPI_Isend(&out,1,MPI_UNSIGNED,2,2,MPI_COMM_WORLD,requests+3);
	MPI_Isend(&b,1,MPI_UNSIGNED,2,1,MPI_COMM_WORLD,requests+4);
	MPI_Isend(&out,1,MPI_UNSIGNED,3,1,MPI_COMM_WORLD,requests+5);
	
	MPI_Waitall(6,requests,statuses);
	
	return out;
}

inline uint32_t round_1_parallel(uint32_t a, uint32_t X_k, uint32_t s, uint32_t T_i, int rank, int num_threads, int iter)
{
	uint32_t b, c, d, temp1, temp2, out;
	MPI_Status statuses[3];
	MPI_Request requests[3];
	//int flags[3] = {0,0,0};
	rank += 3*num_threads;
	//printf("Rank %d: i = %d\n",rank % num_threads,iter);
	if(iter > 3)
	{
		//printf("Rank %d: recv b from rank %d, iter = %d\n",rank%num_threads,(rank-1)%num_threads,iter-1);
		MPI_Irecv(&b,1,MPI_UNSIGNED,(rank-1)%num_threads,iter-1,MPI_COMM_WORLD,requests);
		//printf("Rank %d: recv c from rank %d, iter = %d\n",rank%num_threads,(rank-2)%num_threads,iter-2);
		MPI_Irecv(&c,1,MPI_UNSIGNED,(rank-2)%num_threads,iter-2,MPI_COMM_WORLD,requests+1);
		//printf("Rank %d: recv d from rank %d, iter = %d\n",rank%num_threads,(rank-3)%num_threads,iter-3);
		MPI_Irecv(&d,1,MPI_UNSIGNED,(rank-3)%num_threads,iter-3,MPI_COMM_WORLD,requests+2);
	}
	else
	{
		if(rank%num_threads == 0)
		{
			//printf("Rank %d: recv b, iter = %d\n",rank%num_threads,iter-1);
			MPI_Irecv(&b,1,MPI_UNSIGNED,0,iter-1,MPI_COMM_WORLD,requests);
			//printf("Rank %d: recv c, iter = %d\n",rank%num_threads,iter-2);
			MPI_Irecv(&c,1,MPI_UNSIGNED,0,iter-2,MPI_COMM_WORLD,requests+1);
			//printf("Rank %d: recv d, iter = %d\n",rank%num_threads,iter-3);
			MPI_Irecv(&d,1,MPI_UNSIGNED,0,iter-3,MPI_COMM_WORLD,requests+2);
		}
		if(rank%num_threads == 1)
		{
			//printf("Rank %d: recv b from rank %d, iter = %d\n",rank%num_threads,0,3);
			MPI_Irecv(&b,1,MPI_UNSIGNED,0,3,MPI_COMM_WORLD,requests);
			//printf("Rank %d: recv c from rank %d, iter = %d\n",rank%num_threads,0,2);
			MPI_Irecv(&c,1,MPI_UNSIGNED,0,2,MPI_COMM_WORLD,requests+1);
			//printf("Rank %d: recv d from rank %d, iter = %d\n",rank%num_threads,0,1);
			MPI_Irecv(&d,1,MPI_UNSIGNED,0,1,MPI_COMM_WORLD,requests+2);
		}
		if(rank%num_threads == 2)
		{
			//printf("Rank %d: recv b from rank %d, iter = %d\n",rank%num_threads,(rank-1)%num_threads,iter-1);
			MPI_Irecv(&b,1,MPI_UNSIGNED,(rank-1)%num_threads,iter-1,MPI_COMM_WORLD,requests);
			//printf("Rank %d: recv c from rank %d, iter = %d\n",rank%num_threads,0,2);
			MPI_Irecv(&c,1,MPI_UNSIGNED,0,2,MPI_COMM_WORLD,requests+1);
			//printf("Rank %d: recv d from rank %d, iter = %d\n",rank%num_threads,0,1);
			MPI_Irecv(&d,1,MPI_UNSIGNED,0,1,MPI_COMM_WORLD,requests+2);
		}
		if(rank % num_threads == 3)
		{
			//printf("Rank %d: recv b from rank %d, iter = %d\n",rank%num_threads,(rank-1)%num_threads,iter-1);
			MPI_Irecv(&b,1,MPI_UNSIGNED,(rank-1)%num_threads,iter-1,MPI_COMM_WORLD,requests);
			//printf("Rank %d: recv c from rank %d, iter = %d\n",rank%num_threads,(rank-2)%num_threads,iter-2);
			MPI_Irecv(&c,1,MPI_UNSIGNED,(rank-2)%num_threads,iter-2,MPI_COMM_WORLD,requests+1);
			//printf("Rank %d: recv d from rank %d, iter = %d\n",rank%num_threads,0,1);
			MPI_Irecv(&d,1,MPI_UNSIGNED,0,1,MPI_COMM_WORLD,requests+2);
		}
	}
	//printf("Rank %d: X_k = %x, T_i = %x\n",rank%num_threads,X_k,T_i);
	temp1 = a+X_k+T_i;
	temp2 = 32-s;
	
	//printf("Rank %d: waiting\n",rank%num_threads);
	MPI_Waitall(3,requests,statuses);
	//printf("Rank %d: a = %x, b = %x, c = %x, d = %x\n",rank%num_threads,a,b,c,d);
	/*do
	{
		if(!flags[0])
			MPI_Test(requests,flags,statuses);
		if(!flags[1])
			MPI_Test(requests+1,flags+1,statuses+1);
		if(!flags[2])
			MPI_Test(requests+2,flags+2,statuses+2);
		if(rank%num_threads == 1)
			printf("Rank %d: i = %d, (%d,%d,%d)\n",rank%num_threads,iter,flags[0],flags[1],flags[2]);
	}
	while(!(flags[0] && flags[1] && flags[2]));*/
	
	temp1 += F(b,c,d);
	out = b + ((temp1 << s) | (temp1 >> temp2));
	
	//printf("Rank %d: send out to %d, iter = %d\n",rank%num_threads,(rank+1)%num_threads,iter);
	MPI_Isend(&out,1,MPI_UNSIGNED,(rank+1)%num_threads,iter,MPI_COMM_WORLD,requests);
	//printf("Rank %d: send out to %d, iter = %d\n",rank%num_threads,(rank+2)%num_threads,iter);
	MPI_Isend(&out,1,MPI_UNSIGNED,(rank+2)%num_threads,iter,MPI_COMM_WORLD,requests+1);
	//printf("Rank %d: send out to %d, iter = %d\n",rank%num_threads,(rank+3)%num_threads,iter);
	MPI_Isend(&out,1,MPI_UNSIGNED,(rank+3)%num_threads,iter,MPI_COMM_WORLD,requests+2);
	
	return out;
}

inline uint32_t round_2_parallel(uint32_t a, uint32_t X_k, uint32_t s, uint32_t T_i, int rank, int num_threads, int iter)
{
	uint32_t b, c, d, temp1, temp2, out;
	MPI_Status statuses[3];
	MPI_Request requests[3];
	
	rank += 3*num_threads;
	
	MPI_Irecv(&b,1,MPI_UNSIGNED,(rank-1)%num_threads,iter-1,MPI_COMM_WORLD,requests);
	MPI_Irecv(&c,1,MPI_UNSIGNED,(rank-2)%num_threads,iter-2,MPI_COMM_WORLD,requests+1);
	MPI_Irecv(&d,1,MPI_UNSIGNED,(rank-3)%num_threads,iter-3,MPI_COMM_WORLD,requests+2);
	
	temp1 = a+X_k+T_i;
	temp2 = 32-s;
	
	MPI_Waitall(3,requests,statuses);
	
	temp1 += G(b,c,d);
	out = b + ((temp1 << s) | (temp1 >> temp2));
	
	MPI_Isend(&out,1,MPI_UNSIGNED,(rank+1)%num_threads,iter,MPI_COMM_WORLD,requests);
	MPI_Isend(&out,1,MPI_UNSIGNED,(rank+2)%num_threads,iter,MPI_COMM_WORLD,requests+1);
	MPI_Isend(&out,1,MPI_UNSIGNED,(rank+3)%num_threads,iter,MPI_COMM_WORLD,requests+2);
	
	return out;
}

inline uint32_t round_3_parallel(uint32_t a, uint32_t X_k, uint32_t s, uint32_t T_i, int rank, int num_threads, int iter)
{
	uint32_t b, c, d, temp1, temp2, out;
	MPI_Status statuses[3];
	MPI_Request requests[3];
	
	rank += 3*num_threads;
	
	MPI_Irecv(&b,1,MPI_UNSIGNED,(rank-1)%num_threads,iter-1,MPI_COMM_WORLD,requests);
	MPI_Irecv(&c,1,MPI_UNSIGNED,(rank-2)%num_threads,iter-2,MPI_COMM_WORLD,requests+1);
	MPI_Irecv(&d,1,MPI_UNSIGNED,(rank-3)%num_threads,iter-3,MPI_COMM_WORLD,requests+2);
	
	temp1 = a+X_k+T_i;
	temp2 = 32-s;
	
	MPI_Waitall(3,requests,statuses);
	
	temp1 += H(b,c,d);
	out = b + ((temp1 << s) | (temp1 >> temp2));
	
	MPI_Isend(&out,1,MPI_UNSIGNED,(rank+1)%num_threads,iter,MPI_COMM_WORLD,requests);
	MPI_Isend(&out,1,MPI_UNSIGNED,(rank+2)%num_threads,iter,MPI_COMM_WORLD,requests+1);
	MPI_Isend(&out,1,MPI_UNSIGNED,(rank+3)%num_threads,iter,MPI_COMM_WORLD,requests+2);
	
	return out;
}

inline uint32_t round_4_parallel(uint32_t a, uint32_t X_k, uint32_t s, uint32_t T_i, int rank, int num_threads, int iter)
{
	uint32_t b, c, d, temp1, temp2, out;
	MPI_Status statuses[3];
	MPI_Request requests[3];
	
	rank += 3*num_threads;
	
	MPI_Irecv(&b,1,MPI_UNSIGNED,(rank-1)%num_threads,iter-1,MPI_COMM_WORLD,requests);
	MPI_Irecv(&c,1,MPI_UNSIGNED,(rank-2)%num_threads,iter-2,MPI_COMM_WORLD,requests+1);
	MPI_Irecv(&d,1,MPI_UNSIGNED,(rank-3)%num_threads,iter-3,MPI_COMM_WORLD,requests+2);
	
	temp1 = a+X_k+T_i;
	temp2 = 32-s;
	
	MPI_Waitall(3,requests,statuses);
	
	temp1 += I(b,c,d);
	out = b + ((temp1 << s) | (temp1 >> temp2));
	
	MPI_Isend(&out,1,MPI_UNSIGNED,(rank+1)%num_threads,iter,MPI_COMM_WORLD,requests);
	MPI_Isend(&out,1,MPI_UNSIGNED,(rank+2)%num_threads,iter,MPI_COMM_WORLD,requests+1);
	MPI_Isend(&out,1,MPI_UNSIGNED,(rank+3)%num_threads,iter,MPI_COMM_WORLD,requests+2);
	
	return out;
}

void md5_encrypt(uint32_t* message, uint64_t length, uint32_t *output, int rank, int num_threads)
{
	int i,j;
	uint32_t X[16];
	
	uint32_t AA, BB, CC, DD;
	MPI_Status statuses[4];
	MPI_Request requests[4];
	
	uint32_t A = 0x67452301;
	uint32_t B = 0xefcdab89;
	uint32_t C = 0x98badcfe;
	uint32_t D = 0x10325476;
	
	for(j = 0; j < length; j+=16)
	{
		if(rank == 0)
		{
			for(i = 0; i < 16; i++)
			{
				X[i] = message[i+j];
			}
		}
		MPI_Bcast(X,16,MPI_UNSIGNED,0,MPI_COMM_WORLD);
		
		if(j == 0)
		{
		MPI_Bcast(&A,1,MPI_UNSIGNED,0,MPI_COMM_WORLD);
		MPI_Bcast(&B,1,MPI_UNSIGNED,0,MPI_COMM_WORLD);
		MPI_Bcast(&C,1,MPI_UNSIGNED,0,MPI_COMM_WORLD);
		MPI_Bcast(&D,1,MPI_UNSIGNED,0,MPI_COMM_WORLD);
		}
		
		AA = A;
		BB = B;
		CC = C;
		DD = D;
		
		//MPI_Barrier(MPI_COMM_WORLD);
		if(rank == 0)
		{
			A = round_1(A,B,C,D,X[0],7,T[0]);
		}
		for(i = 0; i < 16; i+= 4)
		{
			if(rank > 14) break;
			if(i > 0)
			{
				if(rank == i%num_threads)
				{
					A = round_1_parallel(A,X[i],7,T[i],rank,num_threads,i);
				}
			}
			if(rank == (i+1)%num_threads)
			{
				D = round_1_parallel(D,X[i+1],12,T[i+1],rank,num_threads,i+1);
			}	
			if(rank == (i+2)%num_threads)
			{
				C = round_1_parallel(C,X[i+2],17,T[i+2],rank,num_threads,i+2);
			}
			if(rank == (i+3)%num_threads)
			{
				B = round_1_parallel(B,X[i+3],22,T[i+3],rank,num_threads,i+3);
			}
		}
		
		for(i = 16; i < 32; i+= 4)
		{
			if(rank > 14) break;
			if(rank == (i%num_threads))
				A = round_2_parallel(A,X[i-15],5,T[i],rank,num_threads,i);
			if(rank == ((i+1)%num_threads))
				D = round_2_parallel(D,X[(i-10)%16],9,T[i+1],rank,num_threads,i+1);
			if(rank == ((i+2)%num_threads))
				C = round_2_parallel(C,X[(i-5)%16],14,T[i+2],rank,num_threads,i+2);
			if(rank == ((i+3)%num_threads))
				B = round_2_parallel(B,X[i-16],20,T[i+3],rank,num_threads,i+3);
		}
		
		for(i = 32; i < 48; i+= 4)
		{
			if(rank > 14) break;
			if(rank == (i%num_threads))
				A = round_3_parallel(A,X[(3*i-91)%16],4,T[i],rank,num_threads,i);
			if(rank == ((i+1)%num_threads))
				D = round_3_parallel(D,X[(3*i-88)%16],11,T[i+1],rank,num_threads,i+1);
			if(rank == ((i+2)%num_threads))
				C = round_3_parallel(C,X[(3*i-85)%16],16,T[i+2],rank,num_threads,i+2);
			if(rank == ((i+3)%num_threads))
				B = round_3_parallel(B,X[(3*i-82)%16],23,T[i+3],rank,num_threads,i+3);
		}
		
		for(i = 48; i < 64; i+= 4)
		{
			if(rank > 14) break;
			if(rank == (i%num_threads))
				A = round_4_parallel(A,X[(3*i-96)%16],6,T[i],rank,num_threads,i);
			if(rank == ((i+1)%num_threads))
				D = round_4_parallel(D,X[(3*i-89)%16],10,T[i+1],rank,num_threads,i+1);
			if(rank == ((i+2)%num_threads))
				C = round_4_parallel(C,X[(3*i-82)%16],15,T[i+2],rank,num_threads,i+2);
			if(rank == ((i+3)%num_threads))
				B = round_4_parallel(B,X[(3*i-91)%16],21,T[i+3],rank,num_threads,i+3);
		}
		
		if(rank == 1)
		{
			D += DD;
		}
		else if(rank == 2)
		{
			C += CC;
		}
		else if(rank == 3)
		{
			B += BB;
		}
		
		else if(rank == 0)
		{
			A += AA;
		}
	}
	
	if(rank == 1)
	{
		MPI_Isend(&D,1,MPI_UNSIGNED,0,65,MPI_COMM_WORLD,requests);
	}
	else if(rank == 2)
	{
		MPI_Isend(&C,1,MPI_UNSIGNED,0,66,MPI_COMM_WORLD,requests);
	}
	else if(rank == 3)
	{
		MPI_Isend(&B,1,MPI_UNSIGNED,0,67,MPI_COMM_WORLD,requests);
	}
	
	else if(rank == 0)
	{
		MPI_Irecv(&D,1,MPI_UNSIGNED,1,65,MPI_COMM_WORLD,requests+3);
		MPI_Irecv(&C,1,MPI_UNSIGNED,2,66,MPI_COMM_WORLD,requests+2);
		MPI_Irecv(&B,1,MPI_UNSIGNED,3,67,MPI_COMM_WORLD,requests+1);
		output[0] = swap_bytes(A);
		MPI_Waitall(3,requests+1,statuses+1);
		
		output[1] = swap_bytes(B);
		output[2] = swap_bytes(C);
		output[3] = swap_bytes(D);
	}
}
